# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import datetime
import logging
import os
import io
from concurrent.futures import ProcessPoolExecutor

from absl import (app, flags)

from coin.exchange.base.strategy_logger import parse_file_name
from coin.support.proto_log.logic.util import (run_from_strat_log_archive,
                                               enumerate_strat_log_into_stratinfos)
from coin.proto.coin_strategy_pb2 import StrategyLog
import google.protobuf.json_format as json_format

FLAGS = flags.FLAGS


class PnlStrategy(object):
  def __init__(self):
    self._first_pnl_balance = {}
    self._latest_pnl_balance = {}

  def on_strat_log(self, timestamp, log):
    pb = StrategyLog()
    pb.ParseFromString(log)
    currency = pb.pnl_balance.currency

    if currency not in self._first_pnl_balance:
      self._first_pnl_balance[currency] = pb.pnl_balance
    self._latest_pnl_balance[currency] = pb.pnl_balance

  def calc_pnl(self, file):
    for currency in self._first_pnl_balance.keys():
      first_pnl_balance_json_str = json_format.MessageToJson(self._first_pnl_balance[currency],
                                                             preserving_proto_field_name=True)
      latest_pnl_balance_json_str = json_format.MessageToJson(self._latest_pnl_balance[currency],
                                                              preserving_proto_field_name=True)
      print('first pnl balance: ', first_pnl_balance_json_str, file=file)
      print('latest pnl balance: ', latest_pnl_balance_json_str, file=file)
      pnl = (self._latest_pnl_balance[currency].total - self._first_pnl_balance[currency].total)
      print('pnl:', pnl, currency, file=file)


def launch(root_dir, trading_date, machine, strategy_name):
  start_time = trading_date  # UTC 0 hour
  hours = 24

  pnl_strategy = PnlStrategy()
  run_from_strat_log_archive(pnl_strategy.on_strat_log,
                             start_time,
                             start_time + datetime.timedelta(hours=hours),
                             root_dir,
                             machine,
                             strategy_name)
  output = io.StringIO()
  print(machine, strategy_name, file=output)
  print('================================================================', file=output)
  pnl_strategy.calc_pnl(output)
  print(output.getvalue())
  output.close()


def main(_):
  trading_date = FLAGS.trading_date
  assert trading_date, '--trading_date must be specified.'
  root_dir = os.path.expanduser(FLAGS.root_dir)

  print('Run for %s:\n' % trading_date)
  strat_infos = enumerate_strat_log_into_stratinfos(root_dir, trading_date, parse_file_name)
  trading_date = datetime.datetime.strptime(trading_date, '%Y%m%d')

  with ProcessPoolExecutor(max_workers=12) as executor:
    for strat_info in strat_infos:
      executor.submit(launch, root_dir, trading_date, strat_info.machine, strat_info.strategy_name)


if __name__ == '__main__':
  flags.DEFINE_string('root_dir',
                      '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log',
                      'Strategy log root directory.')

  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
